import numpy as np
import torch

import torch.nn as nn
from torch import optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms, datasets


# 定义神经网络模型（三层神经网络）
# 注释掉的是没有添加激活函数以及批量标准化的的简单全连接层神经网络
class NN(nn.Module):
    def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
        super(NN, self).__init__()
        # self.layer1 = nn.Linear(in_dim, n_hidden_1)
        self.layer1 = nn.Sequential(
            # 简单全连接神经层，添加批标准化，添加ReLu激活函数
            nn.Linear(in_dim, n_hidden_1),
            nn.BatchNorm1d(n_hidden_1), nn.ReLU(True)
        )
        # self.layer2 = nn.Linear(n_hidden_1, n_hidden_2)
        self.layer2 = nn.Sequential(
            # 简单全连接神经层，添加批标准化，添加ReLu激活函数
            nn.Linear(n_hidden_1, n_hidden_2),
            nn.BatchNorm1d(n_hidden_2), nn.ReLU(True)
        )
        # self.layer3 = nn.Linear(n_hidden_2, out_dim)
        self.layer3 = nn.Sequential(
            # 由于第三层是输入层，所以不需要激活函数和批标准化了，只用一个简单的全连接层输出即可
            nn.Linear(n_hidden_2, out_dim)
        )

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        return x


# 对常用参数提前赋值初始化
batch_size = 64
learning_rate = 1e-2
num_epochs = 20


# 使用torchvision.transforms函数将图片转换成数据并且进行标准化
# transforms函数中有多个方法，在这使用了transforms.ToTensor()和transforms.Normalize()方法
# transforms.Compose() 作用就是一次进行一种或多种操作,在这里就是进行了ToTensor，和Normalize两种操作
# transforms.ToTensor() 将图片转换为Tensor类型数据(一维)，并且Tensor中数据都已经标准化(范围0-1)
# transforms.Normalize([参数1], [参数2]) 将图片中每一个tensor数据先减去第一个参数，再除以第二个参数，此处的作用就是将Tensor数据转换到-1~1范围之内
# 此处的图片是黑白的灰度图片，只有一个颜色通道，当处理的图片是彩色图片时会有(R, G, B)三种颜色通道，那时候transforms.Normalize([a, b, c], [d, e, f])得写成这种形式
# data_tf = transforms.Compose(
#     [transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]
# )
# 但教材中给的数据处理存在问题，处理完的tensor是一个二维数据，我暂时不知道怎么将这些数据给拉平放到模型中
def data_tf(x):
    x = np.array(x, dtype='float32') / 255  # 将每个像素点的值都除以255，就能将这些像素点的值投影到0-1之间
    x = (x - 0.5) / 0.5  # 标准化，将数值范围放在 -1~1 之间
    x = x.reshape((-1,))  # 拉平，能够放到模型中
    x = torch.from_numpy(x)  # 转换维Tensor类型数据
    return x


# 数据下载
train_dataset = datasets.MNIST(
    root='./WriteNumber_data', train=True, transform=data_tf, download=True
)
test_dataset = datasets.MNIST(
    root='./WriteNumber_data', train=False, transform=data_tf, download=True
)

# 定义数据迭代器分批载入(使用 torch.utils.data.DataLoader(目标数据, batch_size分割的每份数据大小, shuffle是否要打乱顺序) 函数读取数据)
# DataLoader是一个可迭代对象，可以通过for循环进行访问
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 将模型实例化，定义损失函数和优化器
# 28 * 28 = 784(在这我写784是防止宝宝看到 Num * Num 会误以为传参时必须得是这样的乘积形式)
# 由于图片横向像素数为28纵向像素数是28，所以一张图片的输入就是28*28(784)个，也就是输入维度是784个
# 至于为什么一层输出维度是300第二层输出层是100，这没啥好说的，随便定，看心情定，最后一层输出必须是10，因为只有0-9这十种结果
model = NN(784, 300, 100, 10)
# 多分类交叉熵作为损失函数
criterion = nn.CrossEntropyLoss()
# 梯度下降作为优化器(在这里我添加了一个推动力,可以解决一些局部极小值情况的发生)
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)

# 开始训练模型
losses = []
access = []
eval_losses = []
eval_access = []
for epoch in range(num_epochs):
    train_loss = 0
    train_acc = 0
    model.train()
    for im, label in train_loader:
        # im中有64张图片的像素点信息   torch.Size([64, 784])
        # label中有64个tag对应着这64张图片     torch.Size([64])
        im = Variable(im)
        label = Variable(label)
        # 前向传播
        out = model(im)  # out.size() => torch.Size([64, 10])
        loss = criterion(out, label)
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # 记录累加每一份batch的误差
        train_loss += loss.item()
        # 计算分类的准确率
        _, pred = out.max(1)  # out.max(1) 返回每一行最大的数以及对应的下标   括号内为0时就按列
        num_correct = (pred == label).sum().item()  # 统计出每一batch中预测正确的为多少
        acc = num_correct / im.shape[0]  # 正确率 = 每一批正确的数 / 每一批数据的数量
        train_acc += acc  # 将每一批的成功率都加上求和，下面再除以批数就可以计算出当前轮数的平均正确率

    losses.append(train_loss / len(train_loader))  # 当前这一轮的平均损失值
    access.append(train_acc / len(train_loader))  # 当前这一论的平均成功率

    print('epoch: {}, Train Loss: {:.6f}, Train Acc: {:.6f}'.format(
                                                                    (epoch + 1),
                                                                    train_loss / len(train_loader),
                                                                    train_acc / len(train_loader)
                                                            ))


# 在测试集上检验效果
eval_loss = 0
eval_acc = 0
model.eval()  # 将模型改为预测模式
for im, label in test_loader:
    im = Variable(im, volatile=True)  # volatile=True表示前向传播时不会保留缓存(因为测试集不需要反向传播，所以用不到缓存)
    label = Variable(label)
    out = model(im)
    loss = criterion(out, label)
    # 记录误差
    eval_loss += loss.item()
    # 记录准确率
    _, pred = out.max(1)
    num_correct = (pred == label).sum().item()
    acc = num_correct / im.shape[0]
    eval_acc += acc
eval_losses.append(eval_loss / len(test_loader))
eval_access.append(eval_acc / len(test_loader))

print('Eval Loss: {:.6f}, Eval Acc: {:.6f}'.format(eval_loss / len(test_loader), eval_acc / len(test_loader)))


